from mindspore import nn

'''
重新构建融合网络，目的为将pad_mode变为"valid"
'''
def reconstruct_layers(fused_layers):
    fused_layers_opt = []

    for layer in fused_layers:
        if isinstance(layer, nn.Conv2d):
            fused_layers_opt.append(nn.Conv2d(layer.in_channels, layer.out_channels, layer.kernel_size,stride=layer.stride,
                                pad_mode='valid',weight_init=layer.weight))
        elif isinstance(layer, nn.MaxPool2d):
            fused_layers_opt.append(nn.MaxPool2d(layer.kernel_size,layer.stride,
                                              pad_mode='valid'))

    return fused_layers_opt


'''
获取当前layer的kernel_size, stride, padding等关键参数
'''
def get_layer_param(layer):

    stride = layer.stride
    kernel_size = layer.kernel_size
    p = 0
    if isinstance(stride, tuple):
        stride = stride[0]

    if isinstance(kernel_size, tuple):
        kernel_size = kernel_size[0]

    if hasattr(layer,'pad_mode'):
        if layer.pad_mode.lower() == 'same':
            p = (kernel_size-1)//2

    return kernel_size,stride, p


'''
用深度搜索的方式找出全部网络结构
'''
def dfs_search_layers(net, layers_list):
    if len(net.cells())==0:
        layers_list.append(net)
        return
    for layer in net.cells():
        dfs_search_layers(layer,layers_list)


'''
计算当前层区域映射到下一层的区域
'''
def next_layer_size(x, y, size, stribe):
    if not isinstance(size, tuple):
        size = (size, size)

    if not isinstance(stribe, tuple):
        stribe = (stribe, stribe)

    next_x = (x - size[0]) // stribe[0] + 1
    next_y = (y - size[1]) // stribe[1] + 1

    return next_x, next_y


'''
计算当前层区域映射到的上一层区域
'''
def pre_layer_area(x1, y1, x2, y2, size, stribe):
    if not isinstance(size, tuple):
        size = (size, size)

    if not isinstance(stribe, tuple):
        stribe = (stribe, stribe)

    border_x = [x1 * stribe[0], (x2 - 1) * stribe[0] + size[0]]
    border_y = [y1 * stribe[1], (y2 - 1) * stribe[1] + size[1]]

    return border_x, border_y


'''
计算量统计，卷积网络中计算量主要来自
MAC计算量
C_in * K * K * H_out * W_out + C_in * C_out * H_out * W_out
'''
def cal_workload(area_map,fuse_layers):
    for j, block in enumerate(area_map.values()):
        s = 0
        print('block',j)
        for i,layer in enumerate(fuse_layers):
            if not isinstance(layer,nn.Conv2d):
                continue
            area = block[i]
    #         print(layer)
            print(layer.cls_name,area)
            HW = (area[0][1]-area[0][0])*(area[1][1]-area[1][0])
            C_in = layer.in_channels
            C_out = layer.out_channels
            if isinstance(layer.kernel_size,tuple):
                K = layer.kernel_size[0]
            else:
                K = layer.kernel_size
            print(C_in*K*K*HW + C_in*C_out*HW)
            s += C_in*K*K*HW + C_in*C_out*HW

        print(s)